"""
聊天助手核心模块

实现与LLM的交互、上下文管理和多轮对话功能。
"""

import logging
import uuid
from typing import Optional, List, Dict, Any
from datetime import datetime

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

from .config import Config
from .models import ChatMessage, ChatSession, ChatResponse, MessageRole

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ChatAssistant:
    """
    AI聊天助手类
    
    提供多轮对话功能，管理聊天会话和上下文。
    """
    
    def __init__(self, config: Optional[Config] = None) -> None:
        """
        初始化聊天助手
        
        Args:
            config: 配置实例，如果为None则使用默认配置
        """
        self.config = config or Config.get_default()
        self.config.validate()
        
        # 初始化LangChain ChatOpenAI
        self._init_llm()
        
        # 当前会话
        self.current_session: Optional[ChatSession] = None
        
        # 系统提示词
        self.system_prompt = (
            "你是一个智能AI助手，能够进行自然、有帮助的对话。"
            "请用中文回答用户的问题，保持友好和专业的语调。"
            "如果用户问题不清楚，请礼貌地要求澄清。"
        )
        
        logger.info("ChatAssistant初始化完成")
    
    def _init_llm(self) -> None:
        """初始化LLM客户端"""
        try:
            self.llm = ChatOpenAI(
                openai_api_base=self.config.api_base,
                openai_api_key=self.config.api_key,
                model=self.config.model_name,
                temperature=self.config.temperature,
                max_retries=self.config.max_retries,
                request_timeout=self.config.request_timeout,
                max_tokens=self.config.max_tokens,
                streaming=self.config.streaming
            )
            logger.info(f"LLM客户端初始化成功，模型: {self.config.model_name}")
        except Exception as e:
            logger.error(f"LLM客户端初始化失败: {e}")
            raise
    
    def create_session(self, session_id: Optional[str] = None) -> ChatSession:
        """
        创建新的聊天会话
        
        Args:
            session_id: 会话ID，如果为None则自动生成
            
        Returns:
            ChatSession: 新创建的会话
        """
        if session_id is None:
            session_id = str(uuid.uuid4())
        
        session = ChatSession(
            session_id=session_id,
            messages=[],
            metadata={"created_by": "ChatAssistant"}
        )
        
        # 添加系统消息
        system_message = ChatMessage(
            role=MessageRole.SYSTEM,
            content=self.system_prompt
        )
        session.add_message(system_message)
        
        self.current_session = session
        logger.info(f"创建新会话: {session_id}")
        
        return session
    
    def set_session(self, session: ChatSession) -> None:
        """
        设置当前会话
        
        Args:
            session: 要设置的会话
        """
        self.current_session = session
        logger.info(f"切换到会话: {session.session_id}")
    
    def get_current_session(self) -> Optional[ChatSession]:
        """
        获取当前会话
        
        Returns:
            Optional[ChatSession]: 当前会话，如果没有则返回None
        """
        return self.current_session
    
    def chat(self, user_input: str, create_session_if_none: bool = True) -> ChatResponse:
        """
        进行聊天对话
        
        Args:
            user_input: 用户输入
            create_session_if_none: 如果没有当前会话是否自动创建
            
        Returns:
            ChatResponse: 聊天响应
        """
        try:
            # 检查是否有当前会话
            if self.current_session is None:
                if create_session_if_none:
                    self.create_session()
                else:
                    return ChatResponse.error_response("没有活动会话，请先创建会话")
            
            # 记录用户消息
            user_message = ChatMessage(
                role=MessageRole.USER,
                content=user_input.strip()
            )
            self.current_session.add_message(user_message)
            
            # 准备消息历史
            messages = self._prepare_messages_for_llm()
            
            # 调用LLM
            logger.info(f"发送请求到LLM，消息数量: {len(messages)}")
            response = self.llm.invoke(messages)
            
            # 处理响应
            assistant_content = response.content if hasattr(response, 'content') else str(response)
            
            # 记录助手响应
            assistant_message = ChatMessage(
                role=MessageRole.ASSISTANT,
                content=assistant_content
            )
            self.current_session.add_message(assistant_message)
            
            logger.info(f"收到LLM响应，长度: {len(assistant_content)}")
            
            return ChatResponse.success_response(
                content=assistant_content,
                metadata={
                    "session_id": self.current_session.session_id,
                    "message_count": self.current_session.get_message_count(),
                    "model": self.config.model_name
                }
            )
            
        except Exception as e:
            error_msg = f"聊天请求失败: {str(e)}"
            logger.error(error_msg)
            return ChatResponse.error_response(error_msg)
    
    def _prepare_messages_for_llm(self, max_history: int = 20) -> List[Any]:
        """
        为LLM准备消息格式
        
        Args:
            max_history: 最大历史消息数量
            
        Returns:
            List[Any]: LangChain消息格式列表
        """
        if not self.current_session:
            return []
        
        # 获取最近的消息
        recent_messages = self.current_session.get_recent_messages(max_history)
        
        # 转换为LangChain消息格式
        langchain_messages = []
        for msg in recent_messages:
            if msg.role == MessageRole.SYSTEM:
                langchain_messages.append(SystemMessage(content=msg.content))
            elif msg.role == MessageRole.USER:
                langchain_messages.append(HumanMessage(content=msg.content))
            elif msg.role == MessageRole.ASSISTANT:
                langchain_messages.append(AIMessage(content=msg.content))
        
        return langchain_messages
    
    def clear_session(self) -> bool:
        """
        清空当前会话的消息历史
        
        Returns:
            bool: 是否成功清空
        """
        if self.current_session:
            self.current_session.clear_messages()
            
            # 重新添加系统消息
            system_message = ChatMessage(
                role=MessageRole.SYSTEM,
                content=self.system_prompt
            )
            self.current_session.add_message(system_message)
            
            logger.info(f"清空会话: {self.current_session.session_id}")
            return True
        
        return False
    
    def get_session_info(self) -> Optional[Dict[str, Any]]:
        """
        获取当前会话信息
        
        Returns:
            Optional[Dict[str, Any]]: 会话信息，如果没有当前会话则返回None
        """
        if not self.current_session:
            return None
        
        return {
            "session_id": self.current_session.session_id,
            "message_count": self.current_session.get_message_count(),
            "created_at": self.current_session.created_at.strftime("%Y-%m-%d %H:%M:%S") if self.current_session.created_at else None,
            "updated_at": self.current_session.updated_at.strftime("%Y-%m-%d %H:%M:%S") if self.current_session.updated_at else None,
            "model": self.config.model_name,
            "api_base": self.config.api_base
        }
    
    def set_system_prompt(self, prompt: str) -> None:
        """
        设置系统提示词
        
        Args:
            prompt: 新的系统提示词
        """
        self.system_prompt = prompt
        logger.info("系统提示词已更新")
        
        # 如果有当前会话，更新系统消息
        if self.current_session and self.current_session.messages:
            # 查找并更新第一个系统消息
            for i, msg in enumerate(self.current_session.messages):
                if msg.role == MessageRole.SYSTEM:
                    self.current_session.messages[i].content = prompt
                    self.current_session.updated_at = datetime.now()
                    break
    
    def get_message_history(self, limit: Optional[int] = None) -> List[ChatMessage]:
        """
        获取消息历史
        
        Args:
            limit: 限制返回的消息数量
            
        Returns:
            List[ChatMessage]: 消息历史列表
        """
        if not self.current_session:
            return []
        
        messages = self.current_session.messages.copy()
        
        # 排除系统消息
        user_assistant_messages = [
            msg for msg in messages 
            if msg.role in [MessageRole.USER, MessageRole.ASSISTANT]
        ]
        
        if limit:
            return user_assistant_messages[-limit:]
        
        return user_assistant_messages
